import torch
from torch import nn

from unet.network.module import RRCNNBlock, UpConv
from unet.network.network_type import NetworkType


class R2UNet(NetworkType):
    @classmethod
    def create_model(cls, config, device):
        return R2UNet(img_ch=config["img_ch"], output_ch=config["output_ch"], t=config["t"])

    def __init__(self, img_ch=3, output_ch=1, t=2):
        super(R2UNet, self).__init__()

        self.Maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.Upsample = nn.Upsample(scale_factor=2)

        self.RRCNN1 = RRCNNBlock(ch_in=img_ch, ch_out=64, t=t)

        self.RRCNN2 = RRCNNBlock(ch_in=64, ch_out=128, t=t)

        self.RRCNN3 = RRCNNBlock(ch_in=128, ch_out=256, t=t)

        self.RRCNN4 = RRCNNBlock(ch_in=256, ch_out=512, t=t)

        self.RRCNN5 = RRCNNBlock(ch_in=512, ch_out=1024, t=t)

        self.Up5 = UpConv(ch_in=1024, ch_out=512)
        self.Up_RRCNN5 = RRCNNBlock(ch_in=1024, ch_out=512, t=t)

        self.Up4 = UpConv(ch_in=512, ch_out=256)
        self.Up_RRCNN4 = RRCNNBlock(ch_in=512, ch_out=256, t=t)

        self.Up3 = UpConv(ch_in=256, ch_out=128)
        self.Up_RRCNN3 = RRCNNBlock(ch_in=256, ch_out=128, t=t)

        self.Up2 = UpConv(ch_in=128, ch_out=64)
        self.Up_RRCNN2 = RRCNNBlock(ch_in=128, ch_out=64, t=t)

        self.Conv_1x1 = nn.Conv2d(64, output_ch, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        # encoding path
        x1 = self.RRCNN1(x)

        x2 = self.Maxpool(x1)
        x2 = self.RRCNN2(x2)

        x3 = self.Maxpool(x2)
        x3 = self.RRCNN3(x3)

        x4 = self.Maxpool(x3)
        x4 = self.RRCNN4(x4)

        x5 = self.Maxpool(x4)
        x5 = self.RRCNN5(x5)

        # decoding + concat path
        d5 = self.Up5(x5)
        d5 = torch.cat((x4, d5), dim=1)
        d5 = self.Up_RRCNN5(d5)

        d4 = self.Up4(d5)
        d4 = torch.cat((x3, d4), dim=1)
        d4 = self.Up_RRCNN4(d4)

        d3 = self.Up3(d4)
        d3 = torch.cat((x2, d3), dim=1)
        d3 = self.Up_RRCNN3(d3)

        d2 = self.Up2(d3)
        d2 = torch.cat((x1, d2), dim=1)
        d2 = self.Up_RRCNN2(d2)

        d1 = self.Conv_1x1(d2)

        return d1
